import torch
import torch.nn as nn

import math

from transcoder.models.block import TRANS_BLOCK

class CSISemanticModel(nn.Module):
    def __init__(self, K, Nc, Nt, Q, SNR_dB, **kwargs):
        super().__init__()
        self.K = K
        self.Nt = Nt
        self.Nc = Nc
        self.Q = Q
        self.sigma = 1 / (10**(SNR_dB / 10))
        self.HEmbed = TRANS_BLOCK(2*self.K*self.Nt, 64, 256, 4)
        
    def forward(self, H, H0):
        """
        H: [batch, K, Nc, 1, Nt] 
        """
        gpu = torch.cuda.current_device()
        batch = H.shape[0]
        H_equ = (torch.zeros(batch, self.Nc, self.K, self.Nt) + 0j).cuda(gpu)
        for i in range(self.K):
            H_equ[:,:,i] = (H[:,i] ).reshape(-1, self.Nc, self.Nt)

        H_equ0 = (torch.zeros(batch, self.Nc, self.K, self.Nt) + 0j).cuda(gpu)
        for i in range(self.K):
            H_equ0[:,:,i] = (H0[:,i] ).reshape(-1, self.Nc, self.Nt)
        
        # 计算每个子载波的信道功率
        power = torch.sum(torch.abs(H_equ * H_equ),[2,3]) # 在用户维度 K 和天线维度 Nt 上求和，得到形状为 (batch, Nc) 的信道总功率
        # 归一化信道增益，H * sqrt(Nt) / sqrt(power) 保证归一化后的信道增益在不同天线数量 Nt 下具有一致的功率尺度
        H_equ = H_equ / torch.sqrt(power.reshape(-1, self.Nc, 1, 1)) * math.sqrt(self.Nt)
        H_equ0 = H_equ0 / torch.sqrt(power.reshape(-1, self.Nc, 1, 1)) * math.sqrt(self.Nt)
        
        HH = H_equ0.reshape(-1, self.Nc, self.K*self.Nt) # [batch, Nc, K*Nt]
        HH = torch.cat((torch.real(HH),torch.imag(HH)), 2) # Input Sequence: [batch, Nc, 2*K*Nt]
        # 交互每个子载波之间的信道增益，获得最优信道增益，每个信道增益包括了Nt根天线发送给K个用户的功率，2是实部虚部
        HH = self.HEmbed(HH) # Transformer: [batch, Nc, 2*K*Nt] -> [batch, Nc, 64]
        
        noise = (torch.randn(batch, self.Nc, self.K, self.Q).cuda(gpu) + 1j * torch.randn(batch, self.Nc, self.K, self.Q).cuda(gpu)) \
            / torch.sqrt(power.reshape(-1, self.Nc, 1, 1)) * math.sqrt(self.sigma / 2) * math.sqrt(self.Nt)
            
        return HH, noise, H_equ
        